package com.zzxx.wechat.util;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.Supplier;

@Slf4j
public class MdcCompletableFuture {

    // 使用默认线程池
    public static <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier) {
        return supplyAsync(supplier, Executors.newCachedThreadPool());
    }

    // 使用自定义线程池
    public static <T> CompletableFuture<T> supplyAsync(Supplier<T> supplier, Executor executor) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        final Thread callerThread = Thread.currentThread();
        return CompletableFuture.supplyAsync(() -> {
            try {
                MDC.setContextMap(contextMap);
                return supplier.get();
            } finally {
                if (Thread.currentThread() != callerThread) {
                    MDC.clear();
                }
            }
        }, executor);
    }

    // 使用默认线程池
    public static CompletableFuture<Void> runAsync(Runnable runnable) {
        return runAsync(runnable, Executors.newCachedThreadPool());
    }

    // 使用自定义线程池
    public static CompletableFuture<Void> runAsync(Runnable runnable, Executor executor) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        final Thread callerThread = Thread.currentThread();
        return CompletableFuture.runAsync(() -> {
            try {
                MDC.setContextMap(contextMap);
                runnable.run();
            } finally {
                if (Thread.currentThread() != callerThread) {
                    MDC.clear();
                }
            }
        }, executor);
    }

    // 处理异常并打印日志
    public static <T> CompletableFuture<T> handleException(CompletableFuture<T> future) {
        return future.exceptionally(ex -> {
            logException(ex);
            return null;
        });
    }

    private static void logException(Throwable ex) {
        log.error("异步任务执行异常:", ex);
    }

    public static Void handleExceptionVoid(Throwable throwable) {
        log.error("异步任务执行异常:", throwable);
        return null;
    }
}
